import os
import shutil
from pathlib import Path

def restructure_folders(source_dir, target_dir):
    """
    将源目录结构:
    A/
        images/
            train/
            val/
            test/
        labels/
            train/
            val/
            test/
    
    重组为目标目录结构:
    B/
        train/
            JPEGImages/
            labels/
        val/
            JPEGImages/
            labels/
        test/
            JPEGImages/
            labels/
    
    并生成train.txt、val.txt和test.txt分别包含训练、验证和测试图像的绝对路径
    """
    # 确保目标路径存在
    target_dir = Path(target_dir)
    if not target_dir.exists():
        os.makedirs(target_dir)
    
    source_dir = Path(source_dir)
    
    # 创建新的目录结构
    for split in ['train', 'val', 'test']:
        # 创建主目录
        split_dir = target_dir / split
        jpeg_dir = split_dir / 'JPEGImages'
        labels_dir = split_dir / 'labels'
        
        os.makedirs(jpeg_dir, exist_ok=True)
        os.makedirs(labels_dir, exist_ok=True)
        
        # 复制图片
        source_images = source_dir / 'images' / split
        if source_images.exists():
            for img in source_images.glob('*'):
                shutil.copy2(img, jpeg_dir)
        
        # 复制标签
        source_labels = source_dir / 'labels' / split
        if source_labels.exists():
            for label in source_labels.glob('*'):
                shutil.copy2(label, labels_dir)
    
    # 生成train.txt、val.txt和test.txt
    for split in ['train', 'val', 'test']:
        txt_file = target_dir / f"{split}.txt"
        jpeg_dir = target_dir / split / 'JPEGImages'
        
        with open(txt_file, 'w') as f:
            if jpeg_dir.exists():
                for img in jpeg_dir.glob('*'):
                    f.write(f"{img.absolute()}\n")
    
    print(f"文件夹结构已重组，并已生成train.txt、val.txt和test.txt文件。")

if __name__ == "__main__":
    # 获取用户输入
    source_directory = r"/home/shuai/Documents/paper_dataset/mydataset/my_usv_aug/integrated_dataset_seaships_usv_weak_nousv5"
    target_directory = r"/home/shuai/darknet/data/mydataset"
    
    # 执行重组
    restructure_folders(source_directory, target_directory)